package middleware

import (
	"context"
	"net/http"
	"strings"

	"zwai-app/config"
	"zwai-app/services"
)

type ctxKey string

const userIDKey ctxKey = "userID"

// AuthMiddleware 验证 JWT 并把 userID 写入 Context
func AuthMiddleware(next http.HandlerFunc, cfg config.Config) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		auth := r.Header.Get("Authorization")
		if !strings.HasPrefix(auth, "Bearer ") {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}
		token := strings.TrimPrefix(auth, "Bearer ")
		uid, err := services.ValidateToken(token, cfg)
		if err != nil {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}
		ctx := context.WithValue(r.Context(), userIDKey, uid)
		next(w, r.WithContext(ctx))
	}
}

// FromContext 获取请求 Context 中的 userID
func FromContext(r *http.Request) string {
	if v := r.Context().Value(userIDKey); v != nil {
		return v.(string)
	}
	return ""
}
